import numpy as np
import matplotlib.pyplot as plt
from matplotlib.widgets import Slider, Button, RadioButtons
import matplotlib.patches as patches
from math import atan2, sqrt, cos, sin, pi

# 设置中文显示
plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]
plt.rcParams["axes.unicode_minus"] = False  # 正确显示负号

class RobotArmIK:
    def __init__(self, l1=114, l2=60):
        """
        初始化双连杆机械臂逆运动学分析模型
        :param l1: 第一连杆长度
        :param l2: 第二连杆长度
        """
        self.l1 = l1
        self.l2 = l2
        
        # 工作空间限制
        self.x_min = 50
        self.x_max = 150
        self.z_max = 50  # 高度限制
        
        # 初始目标位置
        self.target_x = 100
        self.target_y = 30
        
        # 解的类型 (0: 肘上, 1: 肘下)
        self.solution_type = 0
        
        # 创建图形和轴
        self.fig, self.ax = plt.subplots(figsize=(10, 8))
        plt.subplots_adjust(left=0.1, bottom=0.4, right=0.85)  # 留出空间放置控件
        
        # 设置机械臂显示轴
        max_reach = self.l1 + self.l2
        self.ax.set_xlim(0, max_reach * 1.1)  # 第一象限
        self.ax.set_ylim(0, max_reach * 1.1)
        self.ax.set_aspect('equal')
        self.ax.grid(True)
        self.ax.set_title('双连杆机械臂逆运动学分析')
        self.ax.set_xlabel('X坐标')
        self.ax.set_ylabel('Y坐标 (高度)')
        
        # 绘制工作空间边界
        self.ax.axvline(x=self.x_min, color='r', linestyle='--', alpha=0.5, label='工作空间边界')
        self.ax.axvline(x=self.x_max, color='r', linestyle='--', alpha=0.5)
        self.ax.axhline(y=self.z_max, color='r', linestyle='--', alpha=0.5)
        
        # 绘制可达工作空间
        self.plot_workspace()
        
        # 绘制底座
        self.base = plt.Circle((0, 0), 5, color='gray')
        self.ax.add_patch(self.base)
        
        # 初始化连杆和末端执行器
        self.link1, = self.ax.plot([], [], 'b-', lw=6)
        self.link2, = self.ax.plot([], [], 'r-', lw=6)
        self.end_effector, = self.ax.plot([], [], 'go', markersize=10)
        
        # 目标点
        self.target_marker, = self.ax.plot([self.target_x], [self.target_y], 'mo', markersize=10, label='目标点')
        
        # 文本显示
        self.angle_text = self.ax.text(10, max_reach*0.9, '', fontsize=10)
        self.position_text = self.ax.text(10, max_reach*0.8, '', fontsize=10)
        self.status_text = self.ax.text(10, max_reach*0.7, '', fontsize=10)
        
        # 创建交互控件
        self._create_widgets()
        
        # 初始计算与更新
        self.calculate_ik()
        self.update_arm()
        
    def plot_workspace(self):
        """绘制机械臂可达工作空间"""
        # 外边界 (L1+L2)
        theta = np.linspace(0, 2*pi, 360)
        x_outer = (self.l1 + self.l2) * np.cos(theta)
        y_outer = (self.l1 + self.l2) * np.sin(theta)
        
        # 内边界 (|L1-L2|)
        inner_radius = abs(self.l1 - self.l2)
        x_inner = inner_radius * np.cos(theta)
        y_inner = inner_radius * np.sin(theta)
        
        # 绘制工作空间
        self.ax.fill_between(x_outer, y_outer, color='lightblue', alpha=0.2, label='可达工作空间')
        self.ax.fill_between(x_inner, y_inner, color='white', alpha=1)
        self.ax.legend()
        
    def _create_widgets(self):
        """创建交互控件"""
        # 目标X坐标滑块
        ax_x = plt.axes([0.1, 0.25, 0.65, 0.03])
        self.slider_x = Slider(ax_x, '目标X坐标', 
                              0, self.l1 + self.l2, 
                              valinit=self.target_x)
        
        # 目标Y坐标滑块
        ax_y = plt.axes([0.1, 0.2, 0.65, 0.03])
        self.slider_y = Slider(ax_y, '目标Y坐标', 
                              0, self.l1 + self.l2, 
                              valinit=self.target_y)
        
        # 解类型选择按钮
        ax_solution = plt.axes([0.75, 0.7, 0.2, 0.2])
        self.solution_buttons = RadioButtons(ax_solution, ('肘上解', '肘下解'))
        
        # 重置按钮
        ax_reset = plt.axes([0.4, 0.05, 0.2, 0.04])
        self.button_reset = Button(ax_reset, '重置目标点')
        
        # 绑定事件处理函数
        self.slider_x.on_changed(self.update_target_x)
        self.slider_y.on_changed(self.update_target_y)
        self.solution_buttons.on_clicked(self.change_solution)
        self.button_reset.on_clicked(self.reset_target)
        
    def calculate_ik(self):
        """
        计算逆运动学解
        返回: (theta1, theta2) 单位为度，或None（无解时）
        """
        x, y = self.target_x, self.target_y
        
        # 检查目标点是否在可达范围内
        d = sqrt(x**2 + y**2)
        max_reach = self.l1 + self.l2
        min_reach = abs(self.l1 - self.l2)
        
        # 检查是否在工作空间内
        in_workspace = (self.x_min <= x <= self.x_max) and (y <= self.z_max) and (x >= 0) and (y >= 0)
        self.in_workspace = in_workspace
        
        # 检查是否在可达范围内
        if d > max_reach or d < min_reach:
            self.theta1 = None
            self.theta2 = None
            self.ik_possible = False
            return False
        
        self.ik_possible = True
        
        # 计算逆运动学
        # 使用余弦定理
        cos_theta2 = (x**2 + y**2 - self.l1**2 - self.l2**2) / (2 * self.l1 * self.l2)
        # 处理数值计算误差
        cos_theta2 = max(min(cos_theta2, 1), -1)
        
        # 两种可能的解
        theta2_rad_1 = -np.arccos(cos_theta2)  # 肘上
        theta2_rad_2 = np.arccos(cos_theta2)   # 肘下
        
        # 计算对应的theta1
        alpha = atan2(y, x)
        beta_1 = atan2(self.l2 * sin(theta2_rad_1), self.l1 + self.l2 * cos(theta2_rad_1))
        beta_2 = atan2(self.l2 * sin(theta2_rad_2), self.l1 + self.l2 * cos(theta2_rad_2))
        
        theta1_rad_1 = alpha - beta_1
        theta1_rad_2 = alpha - beta_2
        
        # 转换为度
        self.theta1_solutions = [np.degrees(theta1_rad_1), np.degrees(theta1_rad_2)]
        self.theta2_solutions = [np.degrees(theta2_rad_1), np.degrees(theta2_rad_2)]
        
        # 根据选择的解类型更新当前角度
        self.theta1 = self.theta1_solutions[self.solution_type]
        self.theta2 = self.theta2_solutions[self.solution_type]
        
        return True
    
    def update_arm(self):
        """更新机械臂显示"""
        if not self.ik_possible:
            # 无解时清空连杆显示
            self.link1.set_data([], [])
            self.link2.set_data([], [])
            self.end_effector.set_data([], [])
            self.status_text.set_text('目标点不可达')
            self.angle_text.set_text('')
            self.position_text.set_text('')
        else:
            # 计算关节位置
            theta1_rad = np.radians(self.theta1)
            theta2_rad = np.radians(self.theta2)
            total_angle = theta1_rad + theta2_rad
            
            x1 = self.l1 * cos(theta1_rad)
            y1 = self.l1 * sin(theta1_rad)
            
            x2 = x1 + self.l2 * cos(total_angle)
            y2 = y1 + self.l2 * sin(total_angle)
            
            # 更新连杆位置
            self.link1.set_data([0, x1], [0, y1])
            self.link2.set_data([x1, x2], [y1, y2])
            self.end_effector.set_data([x2], [y2])
            
            # 更新文本显示
            ws_status = "在工作空间内" if self.in_workspace else "超出工作空间"
            self.angle_text.set_text(f'θ1: {self.theta1:.2f}°, θ2: {self.theta2:.2f}°')
            self.position_text.set_text(f'目标位置: ({self.target_x:.1f}, {self.target_y:.1f})')
            self.status_text.set_text(f'状态: {ws_status}，解类型: {"肘上" if self.solution_type == 0 else "肘下"}')
        
        self.fig.canvas.draw_idle()
    
    # 事件处理函数
    def update_target_x(self, val):
        self.target_x = val
        self.target_marker.set_xdata([val])
        self.calculate_ik()
        self.update_arm()
    
    def update_target_y(self, val):
        self.target_y = val
        self.target_marker.set_ydata([val])
        self.calculate_ik()
        self.update_arm()
    
    def change_solution(self, label):
        if label == '肘上解':
            self.solution_type = 0
        else:
            self.solution_type = 1
        
        if self.ik_possible:
            self.theta1 = self.theta1_solutions[self.solution_type]
            self.theta2 = self.theta2_solutions[self.solution_type]
            self.update_arm()
    
    def reset_target(self, event):
        self.target_x = 100
        self.target_y = 30
        self.slider_x.set_val(self.target_x)
        self.slider_y.set_val(self.target_y)
        self.target_marker.set_data([self.target_x], [self.target_y])
        self.solution_type = 0
        self.solution_buttons.set_active(0)
        self.calculate_ik()
        self.update_arm()

if __name__ == "__main__":
    # 创建机械臂实例，L1=114，L2=60
    arm = RobotArmIK(l1=114, l2=60)
    plt.show()
    